package set

import (
	"encoding/json"
	"sync"

	"golang.org/x/exp/constraints"

	"gitee.com/sillyman/simpleUtil/common/empty"
)

// hashSet 使用 hashmap 实现的集合
// 数据元素是无序的
type hashSet[T constraints.Ordered] struct {
	hashmap map[T]empty.T
	mutex   *sync.RWMutex
	safe    bool
}

// NewHashSet 创建一个实例
func NewHashSet[T constraints.Ordered](initCap int, safe bool) *hashSet[T] {
	hset := hashSet[T]{
		hashmap: make(map[T]empty.T, initCap),
		safe:    safe,
	}
	if safe {
		hset.mutex = new(sync.RWMutex)
	}

	return &hset
}

// Add 添加元素
func (hset *hashSet[T]) Add(elements ...T) {
	hset.maybeLock()
	defer hset.maybeUnlock()
	for i := range elements {
		if _, ok := hset.hashmap[elements[i]]; !ok {
			hset.hashmap[elements[i]] = empty.V
		}
	}
}

// Remove 删除元素
func (hset *hashSet[T]) Remove(elements ...T) {
	hset.maybeLock()
	defer hset.maybeUnlock()
	for i := range elements {
		if _, ok := hset.hashmap[elements[i]]; ok {
			delete(hset.hashmap, elements[i])
		}
	}
}

// Contains 包含某个元素
func (hset *hashSet[T]) Contains(ele T) bool {
	hset.maybeRLock()
	defer hset.maybeRUnlock()
	_, ok := hset.hashmap[ele]
	return ok
}

// Iterator 迭代
func (hset *hashSet[T]) Iterator(fn func(ele T) bool) {
	hset.maybeRLock()
	defer hset.maybeRUnlock()
	for ele := range hset.hashmap {
		if !fn(ele) {
			return
		}
	}
}

// Len 返回元素的数量
func (hset *hashSet[T]) Len() int {
	hset.maybeRLock()
	defer hset.maybeRUnlock()
	return len(hset.hashmap)
}

// Reset 清空容器
func (hset *hashSet[T]) Reset() {
	hset.maybeLock()
	defer hset.maybeUnlock()
	hset.hashmap = make(map[T]empty.T)
}

// Slice 返回slice
func (hset *hashSet[T]) Slice() []T {
	hset.maybeRLock()
	defer hset.maybeRUnlock()
	slice := make([]T, 0, len(hset.hashmap))
	for ele := range hset.hashmap {
		slice = append(slice, ele)
	}
	return slice
}

// Clone 克隆
func (hset *hashSet[T]) Clone() *hashSet[T] {
	newSet := NewHashSet[T](hset.Len()+128, hset.safe)
	newSet.Add(hset.Clone().Slice()...)
	return newSet
}

// Union 并集
func (hset *hashSet[T]) Union(set *hashSet[T]) *hashSet[T] {
	newSet := hset.Clone()
	newSet.Add(set.Slice()...)
	return newSet
}

// Intersect 交集
func (hset *hashSet[T]) Intersect(set *hashSet[T]) *hashSet[T] {
	newSet := NewHashSet[T](hset.Len()+128, hset.safe)
	hset.Iterator(func(ele T) bool {
		if set.Contains(ele) {
			newSet.Add(ele)
		}
		return true
	})

	return newSet
}

// Difference 差集
func (hset *hashSet[T]) Difference(set *hashSet[T]) *hashSet[T] {
	newSet := NewHashSet[T](hset.Len()+128, hset.safe)
	hset.Iterator(func(ele T) bool {
		if !set.Contains(ele) {
			newSet.Add(ele)
		}
		return true
	})

	return newSet
}

// SymmetricDifference 对称差集
func (hset *hashSet[T]) SymmetricDifference(set *hashSet[T]) *hashSet[T] {
	newSet1 := hset.Difference(set)
	newSet2 := set.Difference(hset)
	return newSet1.Union(newSet2)
}

// MarshalText
func (hset *hashSet[T]) MarshalText() ([]byte, error) {
	return json.Marshal(hset.Slice())
}

// UnmarshalText
func (hset *hashSet[T]) UnmarshalText(data []byte) error {
	elements := make([]T, 0, 1024)
	if err := json.Unmarshal(data, &elements); err != nil {
		return err
	}
	hset.Add(elements...)
	return nil
}

// UnmarshalJSON
func (hset *hashSet[T]) UnmarshalJSON(data []byte) error {
	return hset.UnmarshalText(data)
}

func (hset *hashSet[T]) maybeLock() {
	if hset.safe {
		hset.mutex.Lock()
	}
}

func (hset *hashSet[T]) maybeUnlock() {
	if hset.safe {
		hset.mutex.Unlock()
	}
}

func (hset *hashSet[T]) maybeRLock() {
	if hset.safe {
		hset.mutex.RLock()
	}
}

func (hset *hashSet[T]) maybeRUnlock() {
	if hset.safe {
		hset.mutex.RUnlock()
	}
}
